import time
import numpy as np
import torch
import opuslib_next
from main.base import VADProviderBase

class VAD(VADProviderBase):
    def __init__(self):
        """
        初始化Silero VAD模型
        :param threshold: VAD判断阈值
        :param sampling_rate: 采样率
        :param window_size_samples: 窗口大小
        """
        self.model, utils = torch.hub.load(
            repo_or_dir="tools/snakers_vad",
            source="local",
            model='silero_vad',
            force_reload=True
        )
        
        # 数据缓冲区
        self.audio_buffer = bytearray()
        self.client_have_voice_last_time = time.time()
        self.client_voice_stop = True
        self.client_have_voice_last_time = False
        
        self.decoder = opuslib_next.Decoder(16000, 1)
        self.vad_threshold = 0.5
        self.silence_threshold_ms = 1000 # 如果说话停顿比较长，可以把这个值设置大一些
    
    def is_vad(self, audio_data):
        pass
             
    def is_speech(self, audio_data):
        try:
            self.audio_buffer.extend(audio_data)
            
            # 处理缓冲区中的完整帧（每次处理512采样点）
            client_have_voice = False
            
            while len(self.audio_buffer) >= 512 * 2:
                # 提取前512个采样点（1024字节）
                chunk = self.audio_buffer[: 512 * 2]
                self.audio_buffer = self.audio_buffer[512 * 2 :]
                
                # 转换为模型需要的张量格式
                audio_int16 = np.frombuffer(chunk, dtype=np.int16)
                audio_float32 = audio_int16.astype(np.float32) / 32768.0
                audio_tensor = torch.from_numpy(audio_float32)
                
                # 检测语音活动
                with torch.no_grad():
                    speech_prob = self.model(audio_tensor, 16000)
                client_have_voice = speech_prob >= self.vad_threshold

                # 如果之前有声音，但本次没有声音，且与上次有声音的时间查已经超过了静默阈值，则认为已经说完一句话
                if self.client_have_voice_last_time and not client_have_voice:
                    stop_duration = (time.time() * 1000 - self.client_have_voice_last_time)
                    if stop_duration >= self.silence_threshold_ms:
                        self.client_voice_stop = True
                        
                if client_have_voice:
                    self.client_have_voice_last_time = True
                    self.client_have_voice_last_time = time.time() * 1000
                    
            return client_have_voice
        except Exception as e:
            print(f"Error processing audio packet: {e}")
    
    # def test(self):
    #     model = load_silero_vad()
    #     wav = read_audio('./audios/en.mp3')
    #     speech_timestamps = get_speech_timestamps(
    #         wav,
    #         model,
    #         return_seconds=True,  # Return speech timestamps in seconds (default is samples)
    #     )
    #     print(speech_timestamps)
        
    def reset(self):
        """重置VAD状态"""
        # print("重置VAD状态")
        self.audio_buffer = bytearray()
        # self.state = None